/*
 * Copyright 1999-2018 Alibaba Group Holding Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.alibaba.csp.sentinel.slots.block.flow.param;

import java.lang.reflect.Array;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

import com.alibaba.csp.sentinel.cluster.ClusterStateManager;
import com.alibaba.csp.sentinel.cluster.TokenResult;
import com.alibaba.csp.sentinel.cluster.TokenResultStatus;
import com.alibaba.csp.sentinel.cluster.TokenService;
import com.alibaba.csp.sentinel.cluster.client.TokenClientProvider;
import com.alibaba.csp.sentinel.cluster.server.EmbeddedClusterTokenServerProvider;
import com.alibaba.csp.sentinel.log.RecordLog;
import com.alibaba.csp.sentinel.slotchain.ResourceWrapper;
import com.alibaba.csp.sentinel.slots.block.RuleConstant;
import com.alibaba.csp.sentinel.slots.statistic.cache.CacheMap;
import com.alibaba.csp.sentinel.util.TimeUtil;

/**
 * Rule checker for parameter flow control.
 *
 * @author Eric Zhao
 * @since 0.2.0
 */
public final class ParamFlowChecker {

    public static boolean passCheck(ResourceWrapper resourceWrapper, /*@Valid*/ ParamFlowRule rule, /*@Valid*/ int count,
                             Object... args) {
        if (args == null) {//判断参数是否为空
            return true;//直接返回校验通过
        }
        int paramIdx = rule.getParamIdx();//获取规则配置的参数索引坐标
        if (args.length <= paramIdx) {//如果坐标大于参数个数
            return true;//直接返回校验通过
        }
        // Get parameter value.
        Object value = args[paramIdx];//通过坐标获取参数数组对应下标的参数值
        // Assign value with the result of paramFlowKey method
        if (value instanceof ParamFlowArgument) {//
            value = ((ParamFlowArgument) value).paramFlowKey();
        }
        // If value is null, then pass
        if (value == null) {//如果参数值 为null
            return true;//直接返回校验通过
        }//判断规则配置，是否是集群类型，并且是OPS流控类型
        if (rule.isClusterMode() && rule.getGrade() == RuleConstant.FLOW_GRADE_QPS) {
            return passClusterCheck(resourceWrapper, rule, count, value);//是的话，检查整个集群的统计指标来校验规则
        }//不是的话，检查本地机器的统计指标来校验规则
        return passLocalCheck(resourceWrapper, rule, count, value);
    }

    private static boolean passLocalCheck(ResourceWrapper resourceWrapper, ParamFlowRule rule, int count,
                                          Object value) {
        try {
            if (Collection.class.isAssignableFrom(value.getClass())) {//判断传入参数值是否是Collection集合类型
                for (Object param : ((Collection)value)) {//是的话，遍历Collection集合
                    if (!passSingleValueCheck(resourceWrapper, rule, count, param)) {//通过规则校验单个值
                        return false;//直接返回校验不通过
                    }
                }
            } else if (value.getClass().isArray()) {//判断传入参数值，是否是数组类型
                int length = Array.getLength(value);//获取数组长度
                for (int i = 0; i < length; i++) {//遍历数组
                    Object param = Array.get(value, i);
                    if (!passSingleValueCheck(resourceWrapper, rule, count, param)) {//通过规则校验单个值
                        return false;//不通过 直接返回校验不通过
                    }
                }
            } else {//是其他类型，单个值的类型。 //通过规则校验单个值
                return passSingleValueCheck(resourceWrapper, rule, count, value);
            }
        } catch (Throwable e) {
            RecordLog.warn("[ParamFlowChecker] Unexpected error", e);
        }
        return true;//返回校验通过
    }

    static boolean passSingleValueCheck(ResourceWrapper resourceWrapper, ParamFlowRule rule, int acquireCount,
                                        Object value) {
        if (rule.getGrade() == RuleConstant.FLOW_GRADE_QPS) {//规则类型为QOS模式，目前V1.8.0版本控制台只支持此模式配置
            if (rule.getControlBehavior() == RuleConstant.CONTROL_BEHAVIOR_RATE_LIMITER) {//请求速率限制
                return passThrottleLocalCheck(resourceWrapper, rule, acquireCount, value);
            } else {//只要桶中还有令牌，就可以通过，对应的方法
                return passDefaultLocalCheck(resourceWrapper, rule, acquireCount, value);//默认是快速失败效果。
            }
        } else if (rule.getGrade() == RuleConstant.FLOW_GRADE_THREAD) {//为线程数流控模式，此模式，下面代码存在并发bug.
            Set<Object> exclusionItems = rule.getParsedHotItems().keySet();//获取配置的参数项值集合
            //通过资源名从参数统计类paramterMetric中获取配置的索引下标值已经通过的线程数threadCount
            long threadCount = getParameterMetric(resourceWrapper).getThreadCount(rule.getParamIdx(), value);
            if (exclusionItems.contains(value)) {//判断当前传入的参数值value是否在 配置项值集合中
                //如果在 就获取对应配置项值配置的限流阈值。
                int itemThreshold = rule.getParsedHotItems().get(value);
                return ++threadCount <= itemThreshold;//判断已经线程通过数加+1 是否小于配置的限流阈值。如果是则校验通过，如果不是，则校验不通过。
            }//当前传入参数值value不在配置的参数项集合中。
            //则获取整体单机阈值
            long threshold = (long)rule.getCount();
            //判断已通过线程数加+1 是否小于等于单机阈值。如果是，则校验通过。如果不是，则校验失败。
            return ++threadCount <= threshold;
        }
        return true;
    }

    static boolean passDefaultLocalCheck(ResourceWrapper resourceWrapper, ParamFlowRule rule, int acquireCount,
                                         Object value) {
        ParameterMetric metric = getParameterMetric(resourceWrapper);//通过资源名，获取参数统计类metric
        //获取当前规则的令牌计数器，tokenCounters.  记录令牌桶的令牌数量，用于 QPS 限流
        CacheMap<Object, AtomicLong> tokenCounters = metric == null ? null : metric.getRuleTokenCounter(rule);
        //获取当前规则的令牌桶的最后添加时间计数器timeCounters。 记录令牌桶的最后添加时间，用于 QPS 限流
        CacheMap<Object, AtomicLong> timeCounters = metric == null ? null : metric.getRuleTimeCounter(rule);
        if (tokenCounters == null || timeCounters == null) {//其中一个为null
            return true;//直接返回校验通过
        }
        // Calculate max token count (threshold) 获取配置规则中的参数例外项集合
        Set<Object> exclusionItems = rule.getParsedHotItems().keySet();
        long tokenCount = (long)rule.getCount();//获取配置的最大阈值
        if (exclusionItems.contains(value)) {//判断
            tokenCount = rule.getParsedHotItems().get(value);
        }
        if (tokenCount == 0) {//判断配置的阈值，是否为0
            return false; //直接返回校验不通过
        }
        long maxCount = tokenCount + rule.getBurstCount();//配置的最大阈值，加上突然值，为令牌桶的容量。burstCount值一般默认为0
        if (acquireCount > maxCount) {//判断请求数，是否大于配置的最大阈值。这里acquireCount一般都是1
            return false;//直接返回校验失败
        }
        while (true) {
            long currentTime = TimeUtil.currentTimeMillis();//获取当前时间值
            //往timeCounters集合中添加当前执行时间，进入缓存，代表更新令牌桶容量的时间。
            // 如果集合中存在缓存，就返回上次执行时间。 如果不存在，就加入当前时间作为缓存。
            AtomicLong lastAddTokenTime = timeCounters.putIfAbsent(value, new AtomicLong(currentTime));
            if (lastAddTokenTime == null) {//如果上次执行时间为空，说明这次第一次进入。 初始化令牌桶容量
                // Token never added, just replenish the tokens and consume {@code acquireCount} immediately.
                //往tokenCounters缓存次数集合中，添加还剩maxCount - acquireCount计算出的个数请求通过，在配置的统计时间内。
                tokenCounters.putIfAbsent(value, new AtomicLong(maxCount - acquireCount));
                return true;//直接返回校验通过
            }
            // Calculate the time duration since last token was added.  计算当前执行时间与上次执行时间的间隔。
            long passTime = currentTime - lastAddTokenTime.get();
            // A simplified token bucket algorithm that will replenish the tokens only when statistic window has passed.
            if (passTime > rule.getDurationInSec() * 1000) {//判断时间间隔是否大于配置的统计窗口时长
                //补充令牌桶的容量token
                AtomicLong oldQps = tokenCounters.putIfAbsent(value, new AtomicLong(maxCount - acquireCount));
                if (oldQps == null) {//如果是第一次计算qps
                    // Might not be accurate here.  这种情况应该在这里不会出现。
                    lastAddTokenTime.set(currentTime);
                    return true; //直接校验通过。
                } else {
                    long restQps = oldQps.get();//获取上个统计窗口时长内令牌桶属于的容量 restQps值
                    // 每毫秒应该生成的 token = tokenCount / (rule.getDurationInSec() * 1000)
                    // 再 * passTime 即等于应该补充的 token
                    long toAddCount = (passTime * tokenCount) / (rule.getDurationInSec() * 1000);
                    // 补充的 token 不会超过最大值,
                    // 通过应该补充的token值加上桶内剩余的restQps值
                    // 如果大于令牌桶的最大容量，则将令牌桶补充满。
                    // 如果小于令牌桶的最大容量，则将令牌桶的容量设置为 剩余容量restQps加上补充容量toAddCount并减去当前申请流量acquireCount。
                    long newQps = toAddCount + restQps > maxCount ? (maxCount - acquireCount)
                        : (restQps + toAddCount - acquireCount);
                    if (newQps < 0) {//如果计算出的新的容量值，小于0
                        return false;//直接返回校验失败
                    }
                    if (oldQps.compareAndSet(restQps, newQps)) {//通过CAS更新令牌桶容量。 这里可能会ABA问题，但是概览较低
                        lastAddTokenTime.set(currentTime);
                        return true;//直接返回校验通过。
                    }
                    Thread.yield();//CAS修改失败，让出当前cup的调度。
                }
            } else {//在同一个统计窗口时长内
                AtomicLong oldQps = tokenCounters.get(value);//获取当前value值的旧的ops值
                if (oldQps != null) {//不为空
                    long oldQpsValue = oldQps.get();//获取剩余的请求通过数值
                    if (oldQpsValue - acquireCount >= 0) {//剩余可通过的值，减去当前需要消耗的请求数值，并判断是否大于等于0
                        //表明，还没有达到配置的阈值。 通过CAS更新，剩余可通过请求个数。
                        if (oldQps.compareAndSet(oldQpsValue, oldQpsValue - acquireCount)) {
                            return true;//修改成功，直接返回校验通过
                        }
                    } else {//为空，直接返回校验失败
                        return false;
                    }
                }
                Thread.yield();//CAS修改失败，让出当前cup的调度。
            }
        }
    }

    static boolean passThrottleLocalCheck(ResourceWrapper resourceWrapper, ParamFlowRule rule, int acquireCount,
                                          Object value) {
        ParameterMetric metric = getParameterMetric(resourceWrapper);
        CacheMap<Object, AtomicLong> timeRecorderMap = metric == null ? null : metric.getRuleTimeCounter(rule);
        if (timeRecorderMap == null) {
            return true;
        }
        // Calculate max token count (threshold)
        Set<Object> exclusionItems = rule.getParsedHotItems().keySet();
        long tokenCount = (long)rule.getCount();
        if (exclusionItems.contains(value)) {
            tokenCount = rule.getParsedHotItems().get(value);
        }
        if (tokenCount == 0) {
            return false;
        }
        long costTime = Math.round(1.0 * 1000 * acquireCount * rule.getDurationInSec() / tokenCount);
        while (true) {
            long currentTime = TimeUtil.currentTimeMillis();
            AtomicLong timeRecorder = timeRecorderMap.putIfAbsent(value, new AtomicLong(currentTime));
            if (timeRecorder == null) {
                return true;
            }
            //AtomicLong timeRecorder = timeRecorderMap.get(value);
            long lastPassTime = timeRecorder.get();
            long expectedTime = lastPassTime + costTime;
            if (expectedTime <= currentTime || expectedTime - currentTime < rule.getMaxQueueingTimeMs()) {
                AtomicLong lastPastTimeRef = timeRecorderMap.get(value);
                if (lastPastTimeRef.compareAndSet(lastPassTime, currentTime)) {
                    long waitTime = expectedTime - currentTime;
                    if (waitTime > 0) {
                        lastPastTimeRef.set(expectedTime);
                        try {
                            TimeUnit.MILLISECONDS.sleep(waitTime);
                        } catch (InterruptedException e) {
                            RecordLog.warn("passThrottleLocalCheck: wait interrupted", e);
                        }
                    }
                    return true;
                } else {
                    Thread.yield();
                }
            } else {
                return false;
            }
        }
    }

    private static ParameterMetric getParameterMetric(ResourceWrapper resourceWrapper) {
        // Should not be null.
        return ParameterMetricStorage.getParamMetric(resourceWrapper);
    }

    @SuppressWarnings("unchecked")
    private static Collection<Object> toCollection(Object value) {
        if (value instanceof Collection) {
            return (Collection<Object>)value;
        } else if (value.getClass().isArray()) {
            List<Object> params = new ArrayList<Object>();
            int length = Array.getLength(value);
            for (int i = 0; i < length; i++) {
                Object param = Array.get(value, i);
                params.add(param);
            }
            return params;
        } else {
            return Collections.singletonList(value);
        }
    }

    private static boolean passClusterCheck(ResourceWrapper resourceWrapper, ParamFlowRule rule, int count,
                                            Object value) {
        try {
            Collection<Object> params = toCollection(value);

            TokenService clusterService = pickClusterService();
            if (clusterService == null) {
                // No available cluster client or server, fallback to local or
                // pass in need.
                return fallbackToLocalOrPass(resourceWrapper, rule, count, params);
            }

            TokenResult result = clusterService.requestParamToken(rule.getClusterConfig().getFlowId(), count, params);
            switch (result.getStatus()) {
                case TokenResultStatus.OK:
                    return true;
                case TokenResultStatus.BLOCKED:
                    return false;
                default:
                    return fallbackToLocalOrPass(resourceWrapper, rule, count, params);
            }
        } catch (Throwable ex) {
            RecordLog.warn("[ParamFlowChecker] Request cluster token for parameter unexpected failed", ex);
            return fallbackToLocalOrPass(resourceWrapper, rule, count, value);
        }
    }

    private static boolean fallbackToLocalOrPass(ResourceWrapper resourceWrapper, ParamFlowRule rule, int count,
                                                 Object value) {
        if (rule.getClusterConfig().isFallbackToLocalWhenFail()) {
            return passLocalCheck(resourceWrapper, rule, count, value);
        } else {
            // The rule won't be activated, just pass.
            return true;
        }
    }

    private static TokenService pickClusterService() {
        if (ClusterStateManager.isClient()) {
            return TokenClientProvider.getClient();
        }
        if (ClusterStateManager.isServer()) {
            return EmbeddedClusterTokenServerProvider.getServer();
        }
        return null;
    }

    private ParamFlowChecker() {
    }
}
